import torch.nn as nn
import torch

# class PixelShuffleDecoder(nn.Module):
#     """ Pixel shuffle decoder. """
#     def __init__(self, input_feat_dim=128, num_upsample=2, output_channel=2):
#         super(PixelShuffleDecoder, self).__init__()
#         # Get channel parameters
#         self.channel_conf = self.get_channel_conf(num_upsample)
#
#         # Define the pixel shuffle
#         self.pixshuffle = nn.PixelShuffle(2)
#
#         # Process the feature
#         self.conv_block_lst = []
#         # The input block
#         self.conv_block_lst.append(
#             nn.Sequential(
#                 nn.Conv2d(input_feat_dim, self.channel_conf[0],
#                           kernel_size=3, stride=1, padding=1),
#                 nn.BatchNorm2d(self.channel_conf[0]),
#                 nn.ReLU(inplace=True)
#         ))
#
#         # Intermediate block
#         for channel in self.channel_conf[1:-1]:
#             self.conv_block_lst.append(
#                 nn.Sequential(
#                     nn.Conv2d(channel, channel, kernel_size=3,
#                               stride=1, padding=1),
#                     nn.BatchNorm2d(channel),
#                     nn.ReLU(inplace=True)
#             ))
#
#         # Output block
#         self.conv_block_lst.append(
#             nn.Conv2d(self.channel_conf[-1], output_channel,
#                       kernel_size=1, stride=1, padding=0)
#         )
#         self.conv_block_lst = nn.ModuleList(self.conv_block_lst)
#
#     # Get num of channels based on number of upsampling.
#     def get_channel_conf(self, num_upsample):
#         if num_upsample == 2:
#             return [256, 64, 16]
#         elif num_upsample == 3:
#             return [256, 64, 16, 4]
#
#     def forward(self, input_features):
#         # Iterate til output block
#         out = input_features
#         for block in self.conv_block_lst[:-1]:
#             out = block(out)
#             out = self.pixshuffle(out)
#
#         # Output layer
#         out = self.conv_block_lst[-1](out)
#
#         return out

class PixelShuffleDecoder(nn.Module):
    """ Pixel shuffle decoder. """

    # 对于HourGlass backbone，input_feat_dim=256
    def __init__(self, input_feat_dim=128, num_upsample=2, output_channel=2):
        super(PixelShuffleDecoder, self).__init__()
        # Get channel parameters
        self.channel_conf = self.get_channel_conf(num_upsample)  # [256, 64, 16]

        # Define the pixel shuffle
        self.pixshuffle = nn.PixelShuffle(2)  # 空间放大2倍，通道缩小为原来的1/4

        # Process the feature
        self.conv_block_lst = []
        # The input block
        self.conv_block_lst.append(
            nn.Sequential(
                nn.Conv2d(input_feat_dim, self.channel_conf[0],  # 256
                          kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(self.channel_conf[0]),  # 256
                nn.ReLU(inplace=True)
            ))

        # Intermediate block
        for channel in self.channel_conf[1:-1]:  # 64
            self.conv_block_lst.append(
                nn.Sequential(
                    nn.Conv2d(channel, channel, kernel_size=3,
                              stride=1, padding=1),
                    nn.BatchNorm2d(channel),
                    nn.ReLU(inplace=True)
                ))

        # Output block
        self.conv_block_lst.append(
            nn.Conv2d(self.channel_conf[-1], output_channel,  # 16 2
                      kernel_size=1, stride=1, padding=0)
        )
        self.conv_block_lst = nn.ModuleList(self.conv_block_lst)

        self.combine = nn.Sequential(
            nn.Conv2d(2, 2, kernel_size=1),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True),
            nn.Conv2d(2, 2, kernel_size=1),
            nn.ReLU(inplace=True),
        )

    # Get num of channels based on number of upsampling.
    def get_channel_conf(self, num_upsample):
        if num_upsample == 2:
            return [256, 64, 16]
        elif num_upsample == 3:
            return [256, 64, 16, 4]



    def forward(self, input_features, angle):
        # Iterate til output block
        out = input_features
        # [:-1]从第1项到倒数第二项
        for block in self.conv_block_lst[:-1]:
            out = block(out)
            out = self.pixshuffle(out)

        # Output layer
        out = self.conv_block_lst[-1](out)

        q = torch.zeros_like(out)
        q[:, :1, :, :] = out[:, :1, :, :] - angle
        q[:, 1:, :, :] = out[:, 1:, :, :] + angle
        m = self.combine(q)

        return m
